import numpy as np


def multi_label_precision_recall(inputs:list, targets:list, name_dict:dict=None, eps=0.00001):
    """多标签分类计算PR

    Args:
        inputs (list): 输入预测结果列表
        targets (list): 对应的标签列表
        name_dict (dict, optional): 类别索引映射为字符. Defaults to None.
        eps (float, optional): 精度控制. Defaults to 0.00001.

    Returns:
        tuple: 准确率和每个类别的PR和数量
    """
    targets_set = set(targets)

    cob = np.array([inputs, targets], dtype="<U21").transpose(1, 0)
    right_index = np.where(cob[:, 0] == cob[:, 1])
    ytrue = cob[right_index]

    idict = {"id": [], "precision": [], "recall": [], "count": []}
    for tidx in targets_set:
        tidx = str(tidx)
        TruePos = (ytrue[:, 1] == tidx).sum()
        ToPos = (cob[:, 0] == tidx).sum()
        POS = (cob[:, 1] == tidx).sum()
        if name_dict is not None:
            tidx = name_dict[tidx]
        p = round(TruePos / (ToPos + eps), 2)
        r = round(TruePos / (POS + eps), 2)
        
        idict['id'].append(tidx)
        idict['precision'].append(p)
        idict['recall'].append(r)
        idict['count'].append(POS)
    
    acc = round(ytrue.shape[0] / (cob.shape[0] + eps), 2)
    return acc, idict

def single_label_precision_recall(inputs:list, targets:list, nsplit=10, eps=0.00001):
    """单标签PR曲线数据

    Args:
        inputs (list): 输入预测标签列表
        targets (list): 输入目标标签列表
        nsplit (int, optional): 切分阈值数量. Defaults to 10.
        eps (float, optional): 精度控制. Defaults to 0.00001.

    Returns:
        tuple: 正样本的总数和PR数据字典
    """
    targets_set = set(targets)
    assert len(targets_set) == 2

    cob = np.array([inputs, targets], dtype=np.int).transpose(1, 0)
    cob = np.sort(cob, axis=0)
    pos = cob[:, 1].sum()

    idict = {"threshold": [], "precision": [], "recall": []}
    thresholds = np.linspace(0, 1, nsplit+1)
    for thr in thresholds[:-1]:
        idx = np.where(cob[:, 0] >= thr)
        pos_to_pos = cob[:,1][idx].sum()
        to_pos = idx.shape[0]
        p = round(pos_to_pos / (to_pos + eps), 2)
        r = round(pos_to_pos / (pos + eps), 2)

        idict['threshold'].append(round(thr, 2))
        idict['precision'].append(p)
        idict['recall'].append(r)
    
    return pos, idict


def multi_label_badcase(inputs:list, targets:list, paths:list, name_dict:dict=None):

    targets_set = set(targets)

    cob = np.array([inputs, targets, paths], dtype='<U21').transpose(1, 0)
    wrong_index = np.where(cob[:, 0] != cob[:, 1])
    wrong = cob[wrong_index]

    badcase = {}
    for tidx in targets_set:
        # 标签为类A且预测为非A的索引 Pos ==> Net
        FPIdx = np.where((wrong[:, 0] == str(tidx)) & (wrong[:, 1] != str(tidx)))
        FP = wrong[FPIdx][:, 2]

        # 标签不为类A预测为A的索引 Neg ==> Pos
        FNIdx = np.where((wrong[:, 0] != str(tidx)) & (wrong[:, 1] == str(tidx)))
        FN = wrong[FNIdx][:, 2]

        key = name_dict[tidx] if name_dict else tidx
        badcase[key] = {"FP": FP, "FN": FN}

    return badcase


if __name__ == "__main__":
    a = [3, 1, 3]
    b = [1, 2, 3]
    c = ['a.jpg', 'b.jpg', 'c.jpg']
    name_dict = {1:"哈士奇", 2:"长毛犬", 3:"李德波"}

    badcase = multi_label_badcase(a, b, c, name_dict)
    print(badcase)